// SPDX-License-Identifier: GPL-2.0-only
/*
* Copyright (c) Huawei Technologies Co., Ltd. 2023-2023 All rights reserved.
* Description: Implement sysfs blacklist
* Author: Huawei OS Kernel Lab
* Create: Mon Nov 06 04:28:40 2023
*/

#include <linux/types.h>
#include <linux/list.h>
#include <linux/stringhash.h>
#include <linux/module.h>
#include <linux/printk.h>
#include <linux/errno.h>
#include <linux/rbtree.h>
#include <linux/kernfs.h>
#include <linux/slab.h>
#include <trace/hooks/liblinux.h>

#define PATH_MAX_DEPTH 10 /* Max valid path */

enum name_type {
	NODE_ROOT,
	NODE_LEAF,
	NODE_LEAF_WITH_WHITELIST,
};

struct leaf_name {
	struct list_head list;
	char *name;
	unsigned int len;
};

struct name_hash {
	struct rb_node node;
	struct rb_root root;
	struct list_head leaf_list;
	unsigned int hash;
	unsigned char type;
	unsigned char layer;
};

struct host_path {
	struct name_hash *hit_node;
	unsigned int hash;
};

static struct host_path last_path[PATH_MAX_DEPTH];
static struct rb_root sys_root = RB_ROOT;
static DEFINE_RWLOCK(hashtree_lock);

static unsigned int name2hash(char *name, unsigned int len)
{
	return full_name_hash(NULL, name, len);
}

static struct name_hash *find_name_hash_tree(struct rb_root *root, unsigned int hash)
{
	struct rb_node *rb = root->rb_node;

	while (rb) {
		struct name_hash *tmp;

		tmp = rb_entry(rb, struct name_hash, node);
		if (hash < tmp->hash)
			rb = tmp->node.rb_left;
		else if (hash > tmp->hash)
			rb = tmp->node.rb_right;
		else if (hash == tmp->hash)
			return tmp;
	}

	return NULL;
}

static struct name_hash *create_name_hash_node(struct rb_root *root, unsigned int hash)
{
	struct name_hash *hash_node;
	struct rb_node **link = NULL;
	struct rb_node *parent = NULL;

	hash_node = kmalloc(sizeof(struct name_hash), GFP_ATOMIC);
	if (hash_node == NULL)
		return NULL;

	memset(hash_node, 0, sizeof(struct name_hash));
	INIT_LIST_HEAD(&hash_node->leaf_list);

	link = &root->rb_node;
	while (*link) {
		struct name_hash *tmp;

		parent = *link;
		tmp = rb_entry(parent, struct name_hash, node);
		if (hash < tmp->hash)
			link = &parent->rb_left;
		else
			link = &parent->rb_right;
	}

	rb_link_node(&hash_node->node, parent, link);
	rb_insert_color(&hash_node->node, root);
	return hash_node;
}

static struct name_hash *get_name_hash_tree(struct rb_root *root, unsigned int hash,
					unsigned int layer, unsigned int max_depth)
{
	struct name_hash *hash_node;

	hash_node = find_name_hash_tree(root, hash);
	if (max_depth == layer && hash_node != NULL) { // already had duplicative blacklist
		pr_err("get name_hash, dup path %u\n", layer);
		return ERR_PTR(-EINVAL);
	}

	if (hash_node == NULL) {
		hash_node = create_name_hash_node(root, hash);
		if (hash_node == NULL)
			return ERR_PTR(-ENOMEM);

		hash_node->layer = layer;
		hash_node->hash = hash;
		hash_node->root = RB_ROOT;
	}

	return hash_node;
}

static int parse_path_name(char *path, unsigned int max, unsigned int name_hash[], int *count)
{
	int start_idx = 0;
	int i;
	int len;
	unsigned int depth = 0;

	if (path == NULL)
		return -EINVAL;

	for (i = 0; i < max; i++) {
		unsigned int hash;
		char tmp = path[i];

		if (tmp == '*' || tmp == 0)
			break;
		else if (tmp == '/')
			len = i - start_idx;
		else if (i == max - 1)
			len = max - start_idx;
		else
			continue;

		if (len <= 0 || depth >= PATH_MAX_DEPTH)
			return -EINVAL;

		hash = name2hash(&path[start_idx], (unsigned int)len);
		name_hash[depth] = hash;
		*count = start_idx + len;
		start_idx = i + 1;
		depth++;
	}

	return depth - 1;
}

struct name_hash *add_path_to_hashtree(unsigned int hash[], int depth, struct rb_root *root)
{
	int i = 0;
	struct name_hash *node = NULL;

	for (i = 0; i <= depth; i++) {
		node = get_name_hash_tree(root, hash[i], i, depth);
		if (IS_ERR(node))
			return node;

		/* already had duplicative blacklist */
		if (node->type == NODE_LEAF || node->type == NODE_LEAF_WITH_WHITELIST) {
			pr_err("add path, dup path %d\n", i);
			return ERR_PTR(-EINVAL);
		}
		root = &node->root;
	}

	return node;
}

struct leaf_name *create_leaf_name(char *path, int count)
{
	struct leaf_name *leaf_node;

	leaf_node = kmalloc(sizeof(struct leaf_name), GFP_ATOMIC);
	if (!leaf_node)
		return NULL;

	INIT_LIST_HEAD(&leaf_node->list);
	leaf_node->name = kmalloc(count, GFP_ATOMIC);
	if (!leaf_node->name) {
		kfree(leaf_node);
		return NULL;
	}

	strncpy(leaf_node->name, path, count);
	leaf_node->len = count;
	return leaf_node;
}

static int save_leaf_name(struct name_hash *node, char *path, int count)
{
	struct leaf_name *leaf_node;
	struct leaf_name *tmp;

	if (list_empty(&node->leaf_list)) {
		leaf_node = create_leaf_name(path, count);
		if (leaf_node == NULL)
			return -ENOMEM;

		list_add_tail(&node->leaf_list, &leaf_node->list);
		return 0;
	}

	list_for_each_entry_safe(leaf_node, tmp, &node->leaf_list, list) {
		if (leaf_node->len == count && strncmp(leaf_node->name, path, count))
			return -EINVAL;
	}

	leaf_node = create_leaf_name(path, count);
	if (leaf_node == NULL)
		return -ENOMEM;

	list_add_tail(&node->leaf_list, &leaf_node->list);
	return 0;
}


static int add_blacklist_with_filter(struct name_hash *parent, char *filter[], unsigned int count)
{
	int depth;
	int i;
	unsigned int hash[PATH_MAX_DEPTH] = {0};
	struct name_hash *filter_node = NULL;
	int ret;

	for (i = 0; i < count; i++) {
		int path_len = 0;

		depth = parse_path_name(filter[i], strlen(filter[i]), hash, &path_len);
		if (depth < 0)
			return -EINVAL;

		filter_node = add_path_to_hashtree(hash, depth, &parent->root);
		if (IS_ERR(filter_node))
			return PTR_ERR(filter_node);

		ret = save_leaf_name(filter_node, filter[i], path_len);
		if (ret)
			return ret;

		filter_node->type = NODE_LEAF;
		parent->type = NODE_LEAF_WITH_WHITELIST;
	}

	return 0;
}


int liblinux_add_sysfs_blacklist(char *path, unsigned int len, char *filter[], unsigned int count)
{
	struct rb_root *root = &sys_root;
	unsigned int hash[PATH_MAX_DEPTH] = {0};
	struct name_hash *node = NULL;
	int depth;
	int ret;
	int path_len;

	depth = parse_path_name(path, len, hash, &path_len);
	if (depth < 0)
		return -EINVAL;

	write_lock(&hashtree_lock);
	node = add_path_to_hashtree(hash, depth, root);
	if (IS_ERR(node)) {
		write_unlock(&hashtree_lock);
		return PTR_ERR(node);
	}

	ret = save_leaf_name(node, path, path_len);
	if (ret)
		return ret;

	node->type = NODE_LEAF;

	ret = add_blacklist_with_filter(node, filter, count);
	write_unlock(&hashtree_lock);
	return ret;
}
EXPORT_SYMBOL(liblinux_add_sysfs_blacklist);

static void destroy_namehash_tree(struct rb_root *root)
{
	struct rb_node *next;
	struct name_hash *tmp = NULL;

	next = rb_first(root);
	while (next) {
		struct leaf_name *leaf_node;
		struct leaf_name *leaf_tmp;

		tmp = rb_entry(next, struct name_hash, node);
		if (tmp->root.rb_node != NULL)
			destroy_namehash_tree(&tmp->root);

		next = rb_next(&tmp->node);
		rb_erase(&tmp->node, root);
		list_for_each_entry_safe(leaf_node, leaf_tmp, &tmp->leaf_list, list) {
			list_del(&leaf_node->list);
			kfree(leaf_node->name);
			kfree(leaf_node);
		}
		kfree(tmp);
	}
}

void liblinux_destroy_sysfs_blacklist(void)
{
	struct rb_root *root = &sys_root;

	write_lock(&hashtree_lock);
	destroy_namehash_tree(root);
	memset(last_path, 0, sizeof(struct host_path) * PATH_MAX_DEPTH);
	write_unlock(&hashtree_lock);
}
EXPORT_SYMBOL(liblinux_destroy_sysfs_blacklist);

static struct name_hash *find_host_path(unsigned int namehash, unsigned char layer, bool *hit)
{
	if (layer >= PATH_MAX_DEPTH) {
		pr_err("find host path, Invail path %u\n", layer);
		*hit = false;
		return NULL;
	}

	if (last_path[layer].hash != namehash) {
		*hit = false;
		return NULL;
	}

	*hit = true;
	return last_path[layer].hit_node;
}

static void record_host_path(unsigned int hash, struct name_hash *node, unsigned char layer)
{
	int i;

	for (i = layer; i < PATH_MAX_DEPTH; i++) {
		if (i == layer) {
			last_path[i].hit_node = node;
			last_path[i].hash = hash;
			continue;
		}

		if (last_path[i].hash == 0)
			break;

		last_path[i].hit_node = NULL;
		last_path[i].hash = 0;
	}
}

static struct name_hash *find_name_hash(struct rb_root *root, unsigned int hash,
						unsigned char layer)
{
	struct name_hash *node;
	bool hit;

	node = find_host_path(hash, layer, &hit);
	if (hit == false) {
		node = find_name_hash_tree(root, hash);
		record_host_path(hash, node, layer);
	}

	return node;
}

static bool check_leaf_name(struct name_hash *node, int *icompare, const char *path[],
				int start, int end)
{
	struct leaf_name *leaf_node;
	struct leaf_name *tmp;

	list_for_each_entry_safe(leaf_node, tmp, &node->leaf_list, list) {
		int i;
		int index = *icompare;

		if (index >= leaf_node->len || leaf_node->name == NULL)
			continue;

		for (i = end; i >= start; i--) {
			int count = strlen(path[i]);
			if (index + count > leaf_node->len)
				continue;

			if (strncmp(&leaf_node->name[index], path[i], count) != 0)
				continue;

			index += count + 1; // +1 means skipping "/"
		}

		*icompare = index;
		return true;
	}

	return false;
}

static int rollback_parent_path(struct kernfs_node *parent, const char *name, const char *path[])
{
	struct kernfs_node *tmp = parent;
	int i = 0;

	path[i] = name;
	while (tmp) {
		if (tmp->parent == NULL) // already found /sys node
			break;

		if (i == PATH_MAX_DEPTH - 1)
			return -EINVAL;

		path[++i] = tmp->name;
		tmp = tmp->parent;
	}

	return i;
}

static int find_path_in_blacklist(const char *path[], int depth)
{
	int i;
	struct rb_root *root = &sys_root;
	bool white_flag = false;
	bool hit = false;
	int icompare = 0;
	int iend = depth;

	for (i = depth; i >= 0; i--) {
		unsigned int hash = name2hash(path[i], strlen(path[i]));
		struct name_hash *node = find_name_hash(root, hash, depth - i);

		if (node == NULL)
			break;

		root = &node->root;

		if (node->type == NODE_LEAF && check_leaf_name(node, &icompare, path, i, iend)) {
			hit = true;
			break;
		} else if (node->type == NODE_LEAF_WITH_WHITELIST &&
			check_leaf_name(node, &icompare, path, i, iend)) {
			white_flag = true;
			iend = i - 1;
		}
	}

	hit = white_flag == true ? !hit : hit;
	return hit;
}

static void hook_find_file_in_blacklist(void *args, struct kernfs_node *parent,
					const char *name, bool *hit)
{
	const char *path[PATH_MAX_DEPTH] = {0};
	int depth;

	if (sys_root.rb_node == NULL) {
		*hit = false;
		return;
	}

	depth = rollback_parent_path(parent, name, path);
	if (depth < 0) {
		*hit = false;
		return;
	}

	read_lock(&hashtree_lock);
	*hit = find_path_in_blacklist(path, depth);
	read_unlock(&hashtree_lock);
}
INIT_VENDOR_HOOK(ldk_rvh_find_file_in_blacklist, hook_find_file_in_blacklist);
